import os
import torch
import numpy as np
from NIMBLELayer import NIMBLELayer
from utils import batch_to_tensor_device, save_textured_nimble, smooth_mesh
import pytorch3d.io
from pytorch3d.structures.meshes import Meshes

# -------- config / load same as你的 demo ----------
device = torch.zeros(1).device

pm_dict_name = r"assets/NIMBLE_DICT_9137.pkl"
tex_dict_name = r"assets/NIMBLE_TEX_DICT.pkl"

pm_dict = np.load(pm_dict_name, allow_pickle=True)
pm_dict = batch_to_tensor_device(pm_dict, device)

tex_dict = np.load(tex_dict_name, allow_pickle=True)
tex_dict = batch_to_tensor_device(tex_dict, device)

nimble_mano_vreg = None
if os.path.exists(r"assets/NIMBLE_MANO_VREG.pkl"):
    nimble_mano_vreg = np.load("assets/NIMBLE_MANO_VREG.pkl", allow_pickle=True)
    nimble_mano_vreg = batch_to_tensor_device(nimble_mano_vreg, device)

# --------- create NIMBLE layer: 关闭 PCA，pose_ncomp=60（20 joints × 3） ----------
nlayer = NIMBLELayer(pm_dict, tex_dict, device,
                     use_pose_pca=False,
                     pose_ncomp=60,    # 非 PCA 时传入 60（20*3）
                     shape_ncomp=20,
                     nimble_mano_vreg=nimble_mano_vreg)

# ---- batch size ----
bn = 1

# ---- helper: finger → 对应的 pose bone indices（基于你输出的分析） ----
# 说明：pose_param 中的骨骼 index 是 0..19（共 20）
# 根据你给出的扰动结果，20 个 pose-bone index 的分组大致为：
#  0: root/wrist
#  拇指 bones:  1,2,3      (thumb shorter)
#  食指 bones:  4,5,6,7
#  中指 bones:  8,9,10,11
#  无名指 bones:12,13,14,15
#  小指 bones: 16,17,18,19
thumb_bones  = [1, 2, 3]
index_bones  = [4, 5, 6, 7]
middle_bones = [8, 9, 10, 11]
ring_bones   = [12, 13, 14, 15]
pinky_bones  = [16, 17, 18, 19]

# axis = 0 是你测试中“最有效”的 axis（见 perturb 输出里 best axis 0）
axis = 0

# ---- make two candidate poses: try sign1 and sign2 (有时方向相反，需要翻符号) ----
def make_thumbs_up_pose(nlayer, thumb_sign=-1.0, curl_sign=1.2):
    """
    thumb_sign: 把拇指设为伸直（负或正，依据你机器上显示，若方向不对请换符号）
    curl_sign: 其他手指弯曲的幅度（正号或负号）
    返回 pose_param (bn, nlayer.pose_ncomp)
    """
    pose = torch.zeros(bn, nlayer.pose_ncomp).to(nlayer.device)  # (1, 60)
    # small taper from proximal -> distal (更靠掌的关节需要大一点弯曲)
    # 拇指（3 bones）
    thumb_vals = [thumb_sign * 1.2, thumb_sign * 0.9, thumb_sign * 0.6]
    for j, v in zip(thumb_bones, thumb_vals):
        pose[0, j * 3 + axis] = v

    # 食指/中指/无名指/小指：5~骨链用 4 个骨骼（proximal -> distal）
    curl_vals4 = [curl_sign * 1.4, curl_sign * 1.1, curl_sign * 0.9, curl_sign * 0.6]
    for bone_list in (index_bones, middle_bones, ring_bones, pinky_bones):
        for j, v in zip(bone_list, curl_vals4):
            pose[0, j * 3 + axis] = v

    return pose

# shape / texture defaults
shape_param = torch.zeros(bn, nlayer.shape_ncomp).to(nlayer.device)
tex_param = torch.rand(bn, nlayer.tex_ncomp).to(nlayer.device) - 0.5

# output folder
out = "output_thumbs_up"
os.makedirs(out, exist_ok=True)

# try two sign variants (如果第一个方向朝反，你会在第二个里看到拇指方向相反)
variants = [
    ("thumb_neg_others_pos", make_thumbs_up_pose(nlayer, thumb_sign=-1.0, curl_sign=1.2)),
    ("thumb_pos_others_neg", make_thumbs_up_pose(nlayer, thumb_sign=1.0, curl_sign=-1.2))
]

for name, pose_param in variants:
    # forward
    with torch.no_grad():
        skin_v, muscle_v, bone_v, bone_joints, tex_img = nlayer.forward(pose_param, shape_param, tex_param, handle_collision=True)

    # smooth + save meshes (和你原 demo 相同的处理)
    skin_p3dmesh = Meshes(skin_v, nlayer.skin_f.repeat(bn, 1, 1))
    muscle_p3dmesh = Meshes(muscle_v, nlayer.muscle_f.repeat(bn, 1, 1))
    bone_p3dmesh = Meshes(bone_v, nlayer.bone_f.repeat(bn, 1, 1))

    skin_p3dmesh = smooth_mesh(skin_p3dmesh)
    muscle_p3dmesh = smooth_mesh(muscle_p3dmesh)
    bone_p3dmesh = smooth_mesh(bone_p3dmesh)

    skin_mano_v = nlayer.nimble_to_mano(skin_v, is_surface=True)

    # detach to numpy
    tex_img_np = tex_img.detach().cpu().numpy()
    skin_v_smooth = skin_p3dmesh.verts_padded().detach().cpu().numpy()
    bone_joints_np = bone_joints.detach().cpu().numpy()

    # save
    for i in range(bn):
        np.savetxt(os.path.join(out, f"{name}_{i}_bonej.xyz"), bone_joints_np[i])
        np.savetxt(os.path.join(out, f"{name}_{i}_manov.xyz"), skin_mano_v[i])

        pytorch3d.io.IO().save_mesh(bone_p3dmesh[i], os.path.join(out, f"{name}_{i}_bone.obj"))
        pytorch3d.io.IO().save_mesh(muscle_p3dmesh[i], os.path.join(out, f"{name}_{i}_muscle.obj"))
        save_textured_nimble(os.path.join(out, f"{name}_{i}.obj"), skin_v_smooth[i], tex_img_np[i])

print("保存完成，查看文件夹：", out)

